import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.statespace.sarimax import SARIMAX

# 销售数据
data = {
    '日期': pd.date_range(start='2015-01-31', periods=105, freq='M'),
    '销售量': [
        24211, 23086, 42135, 32252, 37366, 41037, 36529, 37743, 47598, 51698, 57851, 91908,
        35521, 34714, 57998, 49644, 54504, 68652, 62967, 63417, 76876, 63731, 79257, 103527,
        39380, 53066, 90946, 68116, 88816, 100498, 87788, 102889, 122860, 120912, 141032, 173020,
        82035, 80966, 141871, 128457, 159346, 142324, 145475, 172435, 200510, 208820, 237553, 286367,
        153695, 111541, 224335, 166200, 179270, 264591, 148144, 157696, 183393, 149552, 176547, 279214,
        150613, 116170, 192380, 110274, 144600, 229894, 247575, 240981, 345519, 341531, 414368, 571475,
        321031, 269743, 531702, 392498, 442000, 583507, 480506, 516416, 685881, 589663, 721456, 907606,
        603007, 541780, 851489, 542732, 699708, 913479, 778092, 847580, 1040289, 932191, 1058342, 1264645,
        662400, 812487, 1097196, 928739, 1057509, 1260470, 1104592, 1238484, 1291077
    ]
}

# 将数据转换为DataFrame
df = pd.DataFrame(data)

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 绘制销售数据图
plt.figure(figsize=(15, 7))
plt.plot(df['日期'], df['销售量'], marker='o', linestyle='-', color='blue')
plt.title('全球新能源汽车月销售量')
plt.xlabel('日期')
plt.ylabel('销售量（单位）')
plt.grid(True)
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# 训练ARIMA模型
model = ARIMA(df['销售量'], order=(1, 1, 1))
model_fit = model.fit()

# 访问残差
residuals = model_fit.resid

# 绘制残差图
plt.figure(figsize=(10, 4))
plt.plot(df['日期'], residuals, marker='o', linestyle='-', color='blue')
plt.title('ARIMA模型残差')
plt.xlabel('日期')
plt.ylabel('残差')
plt.grid(True)
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# 残差热图
heatmap_data = pd.pivot_table(data=df, values='销售量', index=df['日期'].dt.year, columns=df['日期'].dt.month)
plt.figure(figsize=(10, 6))
sns.heatmap(heatmap_data, cmap='YlGnBu', annot=True, fmt=".0f")
plt.title('ARIMA模型残差热图')
plt.xlabel('月份')
plt.ylabel('年份')
plt.tight_layout()
plt.show()

# SARIMAX模型
sarimax_model_seasonal = SARIMAX(df['销售量'], order=(1,1,1), seasonal_order=(1,1,1,12))
sarimax_model_seasonal_fit = sarimax_model_seasonal.fit(disp=False)

# 进行季节性预测
forecast_seasonal = sarimax_model_seasonal_fit.get_forecast(steps=120)
forecast_index_seasonal = pd.date_range(start=df['日期'].iloc[-1] + pd.DateOffset(months=1), periods=120, freq='M')
forecast_mean_seasonal = forecast_seasonal.predicted_mean
forecast_conf_int_seasonal = forecast_seasonal.conf_int()

# 绘制最终预测结果
plt.figure(figsize=(15, 7))
plt.plot(df['日期'], df['销售量'], label='历史销售量')
plt.plot(forecast_index_seasonal, forecast_mean_seasonal, color='red', label='季节性预测')
plt.fill_between(forecast_index_seasonal, forecast_conf_int_seasonal.iloc[:, 0], forecast_conf_int_seasonal.iloc[:, 1], color='pink', alpha=0.3)
plt.title('未来10年全球新能源汽车月销售量季节性预测')
plt.xlabel('日期')
plt.ylabel('销售量（单位）')
plt.legend()
plt.grid(True)
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()
